#include "SecureDesktopCapture.h"
#include "Utils.h"
#include <iostream>

SecureDesktopCapture::SecureDesktopCapture() :
    m_originalDesktop(NULL),
    m_secureDesktop(NULL),
    m_initialized(false) {
}

SecureDesktopCapture::~SecureDesktopCapture() {
    Cleanup();
}

bool SecureDesktopCapture::Initialize() {
    Utils::DebugLog("SecureDesktopCapture: Initializing...");

    m_originalDesktop = GetThreadDesktop(GetCurrentThreadId());
    if (!m_originalDesktop) {
        m_lastError = "Failed to get original desktop: " + Utils::GetLastWin32Error();
        return false;
    }

    m_initialized = true;
    Utils::DebugLog("SecureDesktopCapture: Initialized successfully");
    return true;
}

void SecureDesktopCapture::Cleanup() {
    if (m_secureDesktop && m_secureDesktop != m_originalDesktop) {
        CloseDesktop(m_secureDesktop);
        m_secureDesktop = NULL;
    }
    m_initialized = false;
}

bool SecureDesktopCapture::CaptureSecureDesktop(HBITMAP& hBitmap, int& width, int& height) {
    if (!m_initialized) {
        m_lastError = "SecureDesktopCapture not initialized";
        return false;
    }

    Utils::DebugLog("SecureDesktopCapture: Attempting secure desktop capture");

    // Method 1: Try to open and switch to Winlogon desktop (where lock screen runs)
    HDESK hWinlogonDesktop = OpenDesktopA("Winlogon", 0, FALSE,
        DESKTOP_READOBJECTS | DESKTOP_WRITEOBJECTS | DESKTOP_CREATEWINDOW | DESKTOP_CREATEMENU);

    bool captureSuccess = false;

    if (hWinlogonDesktop) {
        Utils::DebugLog("SecureDesktopCapture: Successfully opened Winlogon desktop");

        // Save current thread's desktop
        HDESK hCurrentDesktop = GetThreadDesktop(GetCurrentThreadId());

        // Switch to Winlogon desktop
        if (SetThreadDesktop(hWinlogonDesktop)) {
            Utils::DebugLog("SecureDesktopCapture: Successfully switched to Winlogon desktop");

            // Capture from Winlogon desktop
            captureSuccess = CaptureFromCurrentDesktop(hBitmap, width, height);

            // Restore original desktop
            SetThreadDesktop(hCurrentDesktop);
        } else {
            Utils::DebugLog("SecureDesktopCapture: Failed to switch to Winlogon desktop: " + Utils::GetLastWin32Error());
        }

        CloseDesktop(hWinlogonDesktop);
    } else {
        Utils::DebugLog("SecureDesktopCapture: Failed to open Winlogon desktop: " + Utils::GetLastWin32Error());
    }

    // Method 2: Try using GetDC with specific desktop handle
    if (!captureSuccess) {
        Utils::DebugLog("SecureDesktopCapture: Trying alternative method with desktop enumeration");

        // Enumerate all desktops and try each one
        HWINSTA hWinStation = GetProcessWindowStation();
        if (hWinStation) {
            // Try to access current input desktop
            HDESK hInputDesktop = OpenInputDesktop(0, FALSE, DESKTOP_READOBJECTS | DESKTOP_WRITEOBJECTS);
            if (hInputDesktop) {
                HDESK hCurrentDesktop = GetThreadDesktop(GetCurrentThreadId());

                if (SetThreadDesktop(hInputDesktop)) {
                    Utils::DebugLog("SecureDesktopCapture: Switched to input desktop");
                    captureSuccess = CaptureFromCurrentDesktop(hBitmap, width, height);
                    SetThreadDesktop(hCurrentDesktop);
                }

                CloseDesktop(hInputDesktop);
            }
        }
    }

    // Method 3: Try direct screen capture with enhanced permissions
    if (!captureSuccess) {
        Utils::DebugLog("SecureDesktopCapture: Trying direct screen capture method");
        captureSuccess = CaptureDirectScreen(hBitmap, width, height);
    }

    if (!captureSuccess) {
        m_lastError = "All secure desktop capture methods failed";
    }

    return captureSuccess;
}

bool SecureDesktopCapture::SwitchToSecureDesktop() {
    if (!m_secureDesktop) {
        return false;
    }

    if (!SetThreadDesktop(m_secureDesktop)) {
        Utils::DebugLog("SecureDesktopCapture: Failed to switch to secure desktop: " + Utils::GetLastWin32Error());
        return false;
    }

    Utils::DebugLog("SecureDesktopCapture: Successfully switched to secure desktop");
    return true;
}

bool SecureDesktopCapture::CaptureFromCurrentDesktop(HBITMAP& hBitmap, int& width, int& height) {
    int screenWidth = GetSystemMetrics(SM_CXSCREEN);
    int screenHeight = GetSystemMetrics(SM_CYSCREEN);

    if (screenWidth <= 0 || screenHeight <= 0) {
        m_lastError = "Invalid screen dimensions";
        return false;
    }

    // Use GetDC(NULL) to get the screen of the current desktop
    HDC hScreenDC = GetDC(NULL);
    if (!hScreenDC) {
        m_lastError = "Failed to get screen DC: " + Utils::GetLastWin32Error();
        return false;
    }

    HDC hMemDC = CreateCompatibleDC(hScreenDC);
    if (!hMemDC) {
        m_lastError = "Failed to create memory DC: " + Utils::GetLastWin32Error();
        ReleaseDC(NULL, hScreenDC);
        return false;
    }

    HBITMAP hScreenBitmap = CreateCompatibleBitmap(hScreenDC, screenWidth, screenHeight);
    if (!hScreenBitmap) {
        m_lastError = "Failed to create bitmap: " + Utils::GetLastWin32Error();
        DeleteDC(hMemDC);
        ReleaseDC(NULL, hScreenDC);
        return false;
    }

    HGDIOBJ hOldBitmap = SelectObject(hMemDC, hScreenBitmap);

    // Try multiple BitBlt methods
    BOOL result = BitBlt(hMemDC, 0, 0, screenWidth, screenHeight, hScreenDC, 0, 0, SRCCOPY);

    if (!result) {
        // Try with CAPTUREBLT flag for layered windows
        result = BitBlt(hMemDC, 0, 0, screenWidth, screenHeight, hScreenDC, 0, 0, SRCCOPY | CAPTUREBLT);
    }

    SelectObject(hMemDC, hOldBitmap);
    DeleteDC(hMemDC);
    ReleaseDC(NULL, hScreenDC);

    if (result) {
        hBitmap = hScreenBitmap;
        width = screenWidth;
        height = screenHeight;
        Utils::DebugLog("SecureDesktopCapture: Capture from current desktop succeeded");
        return true;
    } else {
        DeleteObject(hScreenBitmap);
        m_lastError = "BitBlt failed: " + Utils::GetLastWin32Error();
        return false;
    }
}

bool SecureDesktopCapture::CaptureDirectScreen(HBITMAP& hBitmap, int& width, int& height) {
    Utils::DebugLog("SecureDesktopCapture: Attempting direct screen capture");

    int screenWidth = GetSystemMetrics(SM_CXSCREEN);
    int screenHeight = GetSystemMetrics(SM_CYSCREEN);

    if (screenWidth <= 0 || screenHeight <= 0) {
        m_lastError = "Invalid screen dimensions for direct capture";
        return false;
    }

    // Create device context for display
    HDC hDisplayDC = CreateDCA("DISPLAY", NULL, NULL, NULL);
    if (!hDisplayDC) {
        m_lastError = "Failed to create display DC: " + Utils::GetLastWin32Error();
        return false;
    }

    HDC hMemDC = CreateCompatibleDC(hDisplayDC);
    if (!hMemDC) {
        m_lastError = "Failed to create memory DC: " + Utils::GetLastWin32Error();
        DeleteDC(hDisplayDC);
        return false;
    }

    HBITMAP hScreenBitmap = CreateCompatibleBitmap(hDisplayDC, screenWidth, screenHeight);
    if (!hScreenBitmap) {
        m_lastError = "Failed to create bitmap: " + Utils::GetLastWin32Error();
        DeleteDC(hMemDC);
        DeleteDC(hDisplayDC);
        return false;
    }

    HGDIOBJ hOldBitmap = SelectObject(hMemDC, hScreenBitmap);

    // Try BitBlt from display DC
    BOOL result = BitBlt(hMemDC, 0, 0, screenWidth, screenHeight, hDisplayDC, 0, 0, SRCCOPY);

    if (!result) {
        // Try StretchBlt as alternative
        result = StretchBlt(hMemDC, 0, 0, screenWidth, screenHeight,
                           hDisplayDC, 0, 0, screenWidth, screenHeight, SRCCOPY);
    }

    SelectObject(hMemDC, hOldBitmap);
    DeleteDC(hMemDC);
    DeleteDC(hDisplayDC);

    if (result) {
        hBitmap = hScreenBitmap;
        width = screenWidth;
        height = screenHeight;
        Utils::DebugLog("SecureDesktopCapture: Direct screen capture succeeded");
        return true;
    } else {
        DeleteObject(hScreenBitmap);
        m_lastError = "Direct screen capture failed: " + Utils::GetLastWin32Error();
        return false;
    }
}

bool SecureDesktopCapture::CaptureFromSecureDesktop(HBITMAP& hBitmap, int& width, int& height) {
    // This method is now deprecated, replaced by CaptureFromCurrentDesktop
    return CaptureFromCurrentDesktop(hBitmap, width, height);
}

bool SecureDesktopCapture::RestoreOriginalDesktop() {
    if (!m_originalDesktop) {
        return false;
    }

    if (!SetThreadDesktop(m_originalDesktop)) {
        Utils::DebugLog("SecureDesktopCapture: Failed to restore original desktop: " + Utils::GetLastWin32Error());
        return false;
    }

    Utils::DebugLog("SecureDesktopCapture: Successfully restored original desktop");
    return true;
}